module FES2DQ412CTE_m
use Mesh2DQ4_m
use FES2DQ4_m
use FES2DQ41_m
use FES2DQ42_m
use Geometry_m
implicit none
type, extends(FES2DQ4) :: FES2DQ412CTE
contains
    procedure :: initFES => initFES2DQ412CTE
    procedure :: assembleHeatToElastic => assembleHeatToElasticQ4CTE
    procedure :: assembleDC1 => assembleDC1_Q4CTE
    procedure :: assembleDC2 => assembleDC2_Q4CTE
    procedure :: extractUVH => extractUVHQ4CTE
    
    procedure :: applyBCZeroDisp => applyBCZeroDispQ4CTE
    procedure :: applyBCHeatQ4CTE_function
    procedure :: applyBCHeatQ4CTE_fromSet
    generic :: applyBCHeat => applyBCHeatQ4CTE_function, applyBCHeatQ4CTE_fromSet
    
end type

interface applyBCHeatQ4CTE
    module procedure applyBCHeatQ4CTE_FromSet
    module procedure applyBCHeatQ4CTE_function
end interface

contains
subroutine initFES2DQ412CTE(this, mesh2D)
    implicit none
    class(FES2DQ412CTE), intent(out) :: this
    type(Mesh2DQ4), target, intent(in) :: mesh2D

    integer :: band
    this%dim = 2
    this%nDof = mesh2D%nv * 3
    this%nVertices = mesh2D%nv
    this%nElement = mesh2D%nt
    this%nBoundary = mesh2D%nb
    this%nPerElement = 4
    this%nPerBoundary = 2
    this%dofPerNode = 3
    this%dofPerElement = 12
    this%dofPerBoundary = 6
    
    band = mesh2D.elementDiff()
    this%nBand = (band + 1) * 3
    this%mesh2D => mesh2D
end subroutine

subroutine assembleHeatToElasticQ4CTE(this, beta, tpCoupled, A)
    implicit none
    class(FES2DQ412CTE), intent(in) :: this
    real*8, intent(in) :: beta
    integer, intent(in) :: tpCoupled
    type(SMatrix) :: A

    integer :: i, j, idx(4), idx2(8)
    real*8 :: xy(4, 2), bbeta(4), HeatToElastic(8, 4)
    bbeta = beta
    do i = 1, this%nElement
        idx = this%mesh2D%elements(i)%idx
        xy(:, 1) = this%mesh2D%nodes(idx)%xy(1)
        xy(:, 2) = this%mesh2D%nodes(idx)%xy(2)
        call getDofFromElement(idx, 2, idx2)
        call elementCoupledMatrixQ4CTE(xy, bbeta, tpCoupled, HeatToElastic)
        call assembleMatrix(A, HeatToElastic, idx2, idx)
    end do
end subroutine

subroutine assembleDC1_Q4CTE(this, rho, c, k, TRef, theta, E, Nu, t, beta, dt, &
                            tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES2DQ412CTE), intent(in) :: this
    real*8, intent(in):: rho, c, k, TRef, theta, E, Nu, t, beta, dt
    type(SMatrix), intent(inout) :: LHS, RHS
    integer, intent(in) :: tpElasticStiff, tpHeatStiff, tpHeatMass, tpCoupled

    integer :: i, j, idx(4), idxHeat(4), idxElastic(8)
    real*8 :: xy(4, 2), coefM(4), coefK(4), coefE(4), coefNu(4), coefBeta(4), coefP(4)
    real*8 :: HeatStiff(4, 4), HeatMass(4, 4)
    real*8 :: ElasticStiff(8, 8), HeatToElastic(8, 4), ElasticToHeat(4, 8)
    type(Mesh2DQ4), pointer :: mesh2D
    mesh2D => this%mesh2D
    coefP = rho
    coefM = rho * c / TRef
    coefK = k / TRef
    coefE = E
    coefNu = Nu
    coefBeta = beta
    do i = 1, mesh2D%nt
        idx = mesh2D%elements(i)%idx
        xy(:, 1) = mesh2D%nodes(idx)%xy(1)
        xy(:, 2) = mesh2D%nodes(idx)%xy(2)
        call getDofFromElement(idx, 2, idxElastic)
        idxHeat = idx + 2 * mesh2D%nv
        call elementElasticityStiff2DQ42(xy, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call elementCoupledMatrixQ4CTE(xy, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call assembleMatrix(LHS, ElasticStiff, idxElastic)
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)

        call elementHeatStiff2DQ41(xy, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass2DQ41(xy, coefM, tpHeatMass, HeatMass)
        call assembleMatrix(LHS, HeatMass + theta * dt * HeatStiff, idxHeat)
        call assembleMatrix(RHS, HeatMass - (1-theta) * dt * HeatStiff, idxHeat)
    end do
end subroutine

subroutine assembleDC2_Q4CTE(this, rho, c, k, TRef, theta, E, Nu, beta, dt, alpha, &
                             tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled, LHS, RHS)
    implicit none
    class(FES2DQ412CTE), intent(in) :: this
    real*8, intent(in):: rho, c, k, TRef, theta, E, Nu, beta, dt, alpha
    integer, intent(in) :: tpElasticStiff, tpElasticMass, tpHeatStiff, tpHeatMass, tpCoupled
    type(SMatrix), intent(inout) :: LHS, RHS

    integer :: i, j, idx(4), idxHeat(4), idxElastic(8)
    real*8 :: xy(4, 2), coefM(4), coefK(4), coefE(4), coefNu(4), coefBeta(4), coefP(4)
    real*8 :: HeatStiff(4, 4), HeatMass(4, 4)
    real*8 :: ElasticStiff(8, 8), ElasticMass(8, 8)
    real*8 :: HeatToElastic(8, 4), ElasticToHeat(4, 8)

    type(Mesh2DQ4), pointer :: mesh2D
    mesh2D => this%mesh2D
    coefP = rho
    coefM = rho * c / TRef
    coefK = k / TRef
    coefE = E
    coefNu = Nu
    coefBeta = beta

    do i = 1, mesh2D%nt
        idx = mesh2D%elements(i)%idx
        xy(:, 1) = mesh2D%nodes(idx)%xy(1)
        xy(:, 2) = mesh2D%nodes(idx)%xy(2)
        call getDofFromElement(idx, 2, idxElastic)
        idxHeat = idx + 2 * mesh2D%nv

        call elementElasticityStiff2DQ42(xy, coefE, coefNu, tpElasticStiff, ElasticStiff)
        call elementElasticityMass2DQ42(xy, coefP, tpElasticMass, ElasticMass)
        call elementCoupledMatrixQ4CTE(xy, coefBeta, tpCoupled, HeatToElastic)
        ElasticToHeat = transpose(HeatToElastic)
        call elementHeatStiff2DQ41(xy, coefK, tpHeatStiff, HeatStiff)
        call elementHeatMass2DQ41(xy, coefM, tpHeatMass, HeatMass)

        call assembleMatrix(LHS, ElasticStiff + 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)     
        call assembleMatrix(LHS, -HeatToElastic, idxElastic, idxHeat)
        call assembleMatrix(LHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(LHS, HeatMass + theta * dt * HeatStiff, idxHeat)

        call assembleMatrix(RHS, 1.0D0 / (alpha * dt * dt) * ElasticMass, idxElastic)
        call assembleMatrix(RHS, ElasticToHeat, idxHeat, idxElastic)
        call assembleMatrix(RHS, HeatMass - (1 - theta) * dt * HeatStiff, idxHeat)    
        
    end do
end subroutine


subroutine elementCoupledMatrixQ4CTE(xy, beta, tpCoupled, K)
    implicit none
    real*8, intent(in) :: xy(4, 2), beta(4)
    integer, intent(in) :: tpCoupled
    real*8, intent(out) :: K(8, 4)

    real*8 :: dx(4), dy(4), u(4), gradient(8)
    real*8 :: KK(8), x(4), y(4)
    real*8 :: mE, mNu, lambda, mu
    real*8 :: xi, et, detJ, wxarea, coefB
    integer :: i, j, si, sj
    
    x = xy(:, 1); y = xy(:, 2)
    K = 0.0D0
    if (tpCoupled == COMMON_2D) then
        do i = 1, 2
            xi = GaussPoints2(i)
            do j = 1, 2
                et = GaussPoints2(j)
                call ISOMap2D(Q4ShapeFunction, 4, x, y, xi, et, u, dx, dy, detJ)
                coefB = dot_product(beta, u)
                gradient(1:8:2) = dx
                gradient(2:8:2) = dy
                wxarea = GaussWeight2(i) * GaussWeight2(j) * detJ * coefB
                do sj = 1, 4
                    K(:, sj) = K(:, sj) + gradient * wxarea * u(sj)  
                end do
            end do
        end do
    else if (tpCoupled == AXISYMMETRIC) then
        do i = 1, 2
            xi = GaussPoints2(i)
            do j = 1, 2
                et = GaussPoints2(j)
                call ISOMap2D(Q4ShapeFunction, 4, x, y, xi, et, u, dx, dy, detJ)
                coefB = dot_product(beta, u)
                gradient(1:8:2) = dx * dot_product(x, u) + u
                gradient(2:8:2) = dy * dot_product(x, u)
                wxarea = GaussWeight2(i) * GaussWeight2(j) * detJ * coefB
                do sj = 1, 4
                    K(:, sj) = K(:, sj) + gradient * wxarea * u(sj)  
                end do
            end do
        end do        
    end if

end subroutine

subroutine extractUVHQ4CTE(this, x, U, V, H)
    implicit none
    class(FES2DQ412CTE), intent(in) :: this
    real*8, intent(in) :: x(this%nDof)
    real*8, intent(out) :: H(this%nDof/3)
    real*8, intent(out) :: U(this%nDof/3)
    real*8, intent(out) :: V(this%nDof/3)

    integer :: nv
    nv = this%nDof/3
    U = x(1:2*nv:2)
    V = x(2:2*nv:2)
    H = x(2*nv+1:this%nDof)
end subroutine

subroutine applyBCZeroDispQ4CTE(this, A, RHS, label, direct)
    use Set_m
    implicit none
    class(FES2DQ412CTE), intent(in) :: this
    type(SMatrix), intent(inout) :: A
    real*8, intent(inout) :: RHS(A%M)
    integer, intent(in) :: label(:)
    integer, intent(in), optional :: direct
    
    type(LinkedSet) :: dirichletBC
    integer, allocatable :: DOF(:)
    type(Mesh2DQ4), pointer :: mesh2D
    mesh2D => this%mesh2D
    call mesh2D.extractBoundaryNodes(dirichletBC, label)
    if (present(direct)) then
        allocate(DOF(dirichletBC%size))
        call getDofFromSet(dirichletBC, 2, DOF, direct)
    else
        allocate(DOF(dirichletBC%size * 2))
        call getDofFromSet(dirichletBC, 2, DOF)
    end if
    call this.applyBCZeroDOF_MR(A, RHS, DOF)
    deallocate(DOF)
    call dirichletBC.empty()
end subroutine

subroutine applyBCHeatQ4CTE_function(this, A, RHS, label, u0, time, tp)
    use Set_m
	implicit none
	class(FES2DQ412CTE), intent(in) :: this
	type(SMatrix), intent(inout) :: A
	real*8, intent(inout) :: RHS(this%nDof)
	integer, intent(in) :: label(:)
	real*8, external :: u0
    real*8, intent(in) :: time
	integer, intent(in) :: tp

	real*8, allocatable :: uu0(:), BCDOF(:)
	integer :: i
	real*8 :: xy(2)
	type(Mesh2DQ4), pointer :: mesh2D
    type(LinkedSet) :: dirichletBC
	type(ele), pointer:: current_p
    mesh2D => this%mesh2D

	call mesh2D.extractBoundaryNodes(dirichletBC, label)

	allocate(uu0(dirichletBC%size))

	i = 0
    current_p => dirichletBC%head
    do while(associated(current_p))
		i = i + 1
		xy = mesh2D%nodes(current_p%n)%xy
		uu0(i) = u0(xy, time)
        current_p%n = current_p%n + 2 * mesh2D%nv
        current_p => current_p%next
    end do
    
    call this.applyBC_SetDOF(A, RHS, dirichletBC, uu0, tp)
	deallocate(uu0)
	call dirichletBC.empty()
end subroutine

subroutine applyBCHeatQ4CTE_FromSet(this, A, RHS, BCSet, BCData, tp)
    use Set_m
	implicit none
	class(FES2DQ412CTE), intent(in) :: this
	type(SMatrix), intent(inout) :: A
	real*8, intent(inout) :: RHS(this%nDof)

    type(LinkedSet), intent(in) :: BCSet
    real*8, intent(in) :: BCData(BCSet%size)
	integer, intent(in) :: tp

    type(LinkedSet) :: myBCSet
	type(ele), pointer:: current_p

    call BCSet.copy(myBCSet)
    call myBCSet.addNumber(2 * this%mesh2D%nv)
	call this.applyBC0(A, RHS, myBCSet, BCData, tp)

	call myBCSet.empty()
end subroutine

end module